# type: ignore
"""
OpenTelemetry 精简配置模块

这个模块提供了 OpenTelemetry 的精简配置，包括：
- 一键式初始化
- 函数追踪装饰器
- 标准日志集成
- FastAPI 和 HTTPX 自动埋点

使用方法：
1. 在主应用中调用 setup_otel() 函数
2. 使用 @trace_function 装饰器自动记录函数调用
3. 使用标准 logger 记录日志（自动集成到 span）
"""

import json
import logging
import functools
from typing import Callable
from datetime import datetime
from opentelemetry import trace
from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.instrumentation.logging import LoggingInstrumentor

# 数据库埋点导入（条件导入，避免依赖问题）
try:
    from opentelemetry.instrumentation.pymysql import PyMySQLInstrumentor

    PYMYSQL_AVAILABLE = True
except ImportError:
    PYMYSQL_AVAILABLE = False

try:
    from opentelemetry.instrumentation.pymongo import PymongoInstrumentor

    PYMONGO_AVAILABLE = True
except ImportError:
    PYMONGO_AVAILABLE = False

from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.aio_pika import AioPikaInstrumentor

# FastAPI 相关导入（用于中间件）
try:
    from fastapi import Request, Response
    from starlette.middleware.base import BaseHTTPMiddleware

    FASTAPI_AVAILABLE = True
except ImportError:
    FASTAPI_AVAILABLE = False

    # 创建占位符类，避免定义错误
    class BaseHTTPMiddleware:
        def __init__(self, app, **kwargs):
            pass

    class Request:
        pass

    class Response:
        pass


def setup_otel(
    service_name: str = 'my-service',
    otlp_endpoint: str = 'http://localhost:4317',
    insecure: bool = True,
    enable_httpx: bool = True,
    app=None,
    enable_trace_id_middleware: bool = True,
    enable_pymysql: bool = True,
    enable_pymongo: bool = True,
    enable_aiohttp: bool = True,
    enable_aio_pika: bool = True,
) -> TracerProvider:
    """
    一键设置 OpenTelemetry

    Args:
        service_name: 服务名称
        otlp_endpoint: OTLP 收集器端点
        insecure: 是否使用不安全连接
        enable_httpx: 是否启用 HTTPX 自动埋点
        app: FastAPI 应用实例（可选）
        enable_trace_id_middleware: 是否自动添加 TraceId 中间件
        enable_pymysql: 是否启用 PyMySQL 自动埋点
        enable_pymongo: 是否启用 PyMongo 自动埋点
        enable_aiohttp: 是否弃用aiohttp自动埋点
        enable_aio_pika: bool = True,  # 是否弃用aio_pika自动埋点

    Returns:
        TracerProvider 实例
    """
    # 创建资源
    resource = Resource(attributes={'service.name': service_name})

    # 创建 TracerProvider
    tracer_provider = TracerProvider(resource=resource)
    trace.set_tracer_provider(tracer_provider)

    # 创建 OTLP 导出器
    otlp_exporter = OTLPSpanExporter(endpoint=otlp_endpoint, insecure=insecure)

    # 添加批量处理器
    tracer_provider.add_span_processor(BatchSpanProcessor(otlp_exporter))

    # 启用自动埋点
    if enable_httpx:
        HTTPXClientInstrumentor().instrument()

    if enable_aiohttp:
        AioHttpClientInstrumentor().instrument()

    if enable_aio_pika:
        AioPikaInstrumentor().instrument()

    # 启用数据库自动埋点
    if enable_pymysql and PYMYSQL_AVAILABLE:
        PyMySQLInstrumentor().instrument()

    if enable_pymongo and PYMONGO_AVAILABLE:
        PymongoInstrumentor().instrument()

    LoggingInstrumentor().instrument(set_logging_format=False)

    # 为 FastAPI 应用添加链路追踪
    if app is not None:
        FastAPIInstrumentor().instrument_app(app, tracer_provider=tracer_provider, exclude_spans=['receive', 'send'])

        # 自动添加 TraceId 中间件
        if enable_trace_id_middleware and FASTAPI_AVAILABLE:
            app.add_middleware(TraceIdMiddleware)

    return tracer_provider


def get_tracer(name: str = __name__):
    """获取 Tracer 实例"""
    return trace.get_tracer(name)


class SmartLogger:
    """智能日志记录器 - 自动集成到 OpenTelemetry span"""

    def __init__(self, name: str = __name__, level: int = logging.INFO):
        self.logger = logging.getLogger(name)
        self.logger.setLevel(level)

        # 如果没有处理器，添加控制台处理器
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)

    def _log_with_span(self, level: str, message: str, *args, **kwargs):
        """将日志同时输出到控制台和当前 span"""
        # 格式化消息
        formatted_message = message % args if args else message

        # 输出到标准日志
        getattr(self.logger, level)(formatted_message)

        # 集成到当前 span
        current_span = trace.get_current_span()
        if current_span and current_span.is_recording():
            event_attributes = {'log.level': level.upper(), 'log.message': formatted_message, 'timestamp': datetime.now().isoformat()}

            # 添加额外属性
            if 'extra' in kwargs:
                event_attributes.update(kwargs['extra'])

            current_span.add_event(f'log.{level}', attributes=event_attributes)

    def info(self, message: str, *args, **kwargs):
        """记录 INFO 级别日志"""
        self._log_with_span('info', message, *args, **kwargs)

    def error(self, message: str, *args, **kwargs):
        """记录 ERROR 级别日志"""
        self._log_with_span('error', message, *args, **kwargs)

    def warning(self, message: str, *args, **kwargs):
        """记录 WARNING 级别日志"""
        self._log_with_span('warning', message, *args, **kwargs)

    def debug(self, message: str, *args, **kwargs):
        """记录 DEBUG 级别日志"""
        self._log_with_span('debug', message, *args, **kwargs)

    def exception(self, message: str, *args, **kwargs):
        """记录异常日志"""
        import sys

        exc_type, exc_value, exc_traceback = sys.exc_info()
        if exc_value:
            kwargs.setdefault('extra', {}).update(
                {'exception.type': exc_type.__name__ if exc_type else None, 'exception.message': str(exc_value)}
            )
        self.error(message, *args, **kwargs)


def get_logger(name: str = __name__) -> SmartLogger:
    """获取智能日志记录器"""
    return SmartLogger(name)


def trace_function(include_args: bool = True, include_result: bool = True, include_execution_time: bool = True):
    """
    函数追踪装饰器

    Args:
        include_args: 是否记录函数参数
        include_result: 是否记录返回结果
        include_execution_time: 是否记录执行时间

    Example:
        @trace_function()
        def my_function(a, b):
            return a + b
    """

    def decorator(func: Callable) -> Callable:
        tracer = trace.get_tracer(__name__)
        logger = SmartLogger(func.__module__)

        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            import time

            function_name = f'{func.__module__}.{func.__qualname__}'
            start_time = time.time()

            with tracer.start_as_current_span(function_name) as span:
                span.set_attribute('function.name', func.__name__)
                span.set_attribute('function.module', func.__module__)

                if include_args:
                    span.set_attribute('function.args.count', len(args))
                    span.set_attribute('function.kwargs.count', len(kwargs))

                    # 记录具体的参数值（限制长度避免过大）
                    for i, arg in enumerate(args):
                        arg_str = str(arg)[:500]  # 限制参数长度
                        span.set_attribute(f'function.args.{i}', arg_str)

                    for key, value in kwargs.items():
                        value_str = str(value)[:500]  # 限制参数长度
                        span.set_attribute(f'function.kwargs.{key}', value_str)

                try:
                    result = await func(*args, **kwargs)
                    execution_time = time.time() - start_time

                    if include_execution_time:
                        span.set_attribute('function.execution_time', execution_time)

                    # 记录返回结果（限制长度避免过大）
                    if include_result:
                        result_str = str(result)[:1000]  # 限制结果长度
                        span.set_attribute('function.result', result_str)

                    # 记录函数调用日志
                    log_data = {
                        'function': function_name,
                        'execution_time': round(execution_time, 4) if include_execution_time else None,
                        'args_count': len(args) if include_args else None,
                        'kwargs_count': len(kwargs) if include_args else None,
                        'result': str(result)[:200] if include_result else None,  # 限制结果长度
                    }

                    logger.info(
                        'Function executed successfully: %s',
                        json.dumps({k: v for k, v in log_data.items() if v is not None}, ensure_ascii=False),
                    )

                    return result

                except Exception as e:
                    execution_time = time.time() - start_time

                    span.set_attribute('error', True)
                    span.set_attribute('error.type', type(e).__name__)
                    span.set_attribute('error.message', str(e))

                    logger.error('Function execution failed: %s - %s', function_name, str(e))
                    raise

        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            import time

            function_name = f'{func.__module__}.{func.__qualname__}'
            start_time = time.time()

            with tracer.start_as_current_span(function_name) as span:
                span.set_attribute('function.name', func.__name__)
                span.set_attribute('function.module', func.__module__)

                if include_args:
                    span.set_attribute('function.args.count', len(args))
                    span.set_attribute('function.kwargs.count', len(kwargs))

                    # 记录具体的参数值（限制长度避免过大）
                    for i, arg in enumerate(args):
                        arg_str = str(arg)[:500]  # 限制参数长度
                        span.set_attribute(f'function.args.{i}', arg_str)

                    for key, value in kwargs.items():
                        value_str = str(value)[:500]  # 限制参数长度
                        span.set_attribute(f'function.kwargs.{key}', value_str)

                try:
                    result = func(*args, **kwargs)
                    execution_time = time.time() - start_time

                    if include_execution_time:
                        span.set_attribute('function.execution_time', execution_time)

                    # 记录返回结果（限制长度避免过大）
                    if include_result:
                        result_str = str(result)[:1000]  # 限制结果长度
                        span.set_attribute('function.result', result_str)

                    # 记录函数调用日志
                    log_data = {
                        'function': function_name,
                        'execution_time': round(execution_time, 4) if include_execution_time else None,
                        'args_count': len(args) if include_args else None,
                        'kwargs_count': len(kwargs) if include_args else None,
                        'result': str(result)[:200] if include_result else None,  # 限制结果长度
                    }

                    logger.info(
                        'Function executed successfully: %s',
                        json.dumps({k: v for k, v in log_data.items() if v is not None}, ensure_ascii=False),
                    )

                    return result

                except Exception as e:
                    execution_time = time.time() - start_time

                    span.set_attribute('error', True)
                    span.set_attribute('error.type', type(e).__name__)
                    span.set_attribute('error.message', str(e))

                    logger.error('Function execution failed: %s - %s', function_name, str(e))
                    raise

        # 根据函数类型返回相应的包装器
        import asyncio

        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper

    return decorator


# FastAPI 中间件
class TraceIdMiddleware(BaseHTTPMiddleware):
    """自动为所有响应添加 traceId 的中间件"""

    def __init__(self, app, **kwargs):
        if not FASTAPI_AVAILABLE:
            raise ImportError('FastAPI is required to use TraceIdMiddleware')
        super().__init__(app, **kwargs)

    async def dispatch(self, request: Request, call_next):
        if request.headers.get('content-type', '').startswith('application/json'):
            body_bytes = await request.body()
            if body_bytes:
                try:
                    body_data = json.loads(body_bytes.decode('utf-8'))
                except json.JSONDecodeError:
                    body_data = None
            else:
                body_data = None
            span = trace.get_current_span()
            if span and span.is_recording():
                span.add_event('request.body', body_data)

            async def receive():
                return {'type': 'http.request', 'body': body_bytes}

            request._receive = receive

        response = await call_next(request)

        # 获取当前的 traceId
        trace_id = get_trace_id()

        # 如果有 traceId，添加到响应头中
        if trace_id:
            response.headers['X-Trace-Id'] = trace_id

            # 如果响应是 JSON 格式，尝试在响应体中也添加 traceId
            if response.headers.get('content-type', '').startswith('application/json'):
                try:
                    # 读取原始响应体
                    body = b''
                    async for chunk in response.body_iterator:
                        body += chunk

                    # 解析 JSON
                    if body:
                        try:
                            json_data = json.loads(body.decode())
                            # 如果是字典，添加 traceId
                            if isinstance(json_data, dict):
                                json_data['traceId'] = trace_id
                                # 重新编码响应体
                                new_body = json.dumps(json_data, ensure_ascii=False).encode()
                                response.headers['content-length'] = str(len(new_body))

                                # 创建新的响应
                                return Response(
                                    content=new_body,
                                    status_code=response.status_code,
                                    headers=dict(response.headers),
                                    media_type=response.media_type,
                                )
                        except (json.JSONDecodeError, UnicodeDecodeError):
                            # 如果解析失败，返回原始响应
                            pass

                    # 如果无法处理，返回原始响应
                    return Response(
                        content=body, status_code=response.status_code, headers=dict(response.headers), media_type=response.media_type
                    )
                except Exception:
                    # 如果处理过程中出错，返回原始响应
                    pass

        return response


# 便捷函数
def get_trace_id() -> str:
    """
    获取当前请求的 traceId

    Returns:
        str: 当前 span 的 traceId，如果没有活跃的 span 则返回空字符串
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        trace_id = current_span.get_span_context().trace_id
        # 将 trace_id 转换为 16 进制字符串格式
        return format(trace_id, '032x')
    return ''


def get_span_id() -> str:
    """
    获取当前请求的 spanId

    Returns:
        str: 当前 span 的 spanId，如果没有活跃的 span 则返回空字符串
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        span_id = current_span.get_span_context().span_id
        # 将 span_id 转换为 16 进制字符串格式
        return format(span_id, '016x')
    return ''


def get_trace_info() -> dict:
    """
    获取当前请求的完整链路追踪信息

    Returns:
        dict: 包含 traceId 和 spanId 的字典
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        span_context = current_span.get_span_context()
        return {
            'traceId': format(span_context.trace_id, '032x'),
            'spanId': format(span_context.span_id, '016x'),
            'traceFlags': span_context.trace_flags,
        }
    return {'traceId': '', 'spanId': '', 'traceFlags': 0}


def log_http_request(method: str, url: str, headers: dict = None):
    """记录HTTP请求"""
    logger = get_logger('http')
    log_data = {'method': method, 'url': url}
    if headers:
        log_data['headers'] = dict(headers)

    logger.info('HTTP Request: %s', json.dumps(log_data, ensure_ascii=False))

    # 添加到当前 span
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        current_span.add_event('http.request', attributes={'http.method': method, 'http.url': url})


def log_http_response(status_code: int, headers: dict = None):
    """记录HTTP响应"""
    logger = get_logger('http')
    log_data = {'status_code': status_code}
    if headers:
        log_data['headers'] = dict(headers)

    logger.info('HTTP Response: %s', json.dumps(log_data, ensure_ascii=False))

    # 添加到当前 span
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        current_span.add_event('http.response', attributes={'http.status_code': status_code})
